#ifndef FACERECOGNIZER_H
#define FACERECOGNIZER_H

#include <dlib/dnn.h>
#include <dlib/image_processing.h>
#include <dlib/image_io.h>
#include <dlib/image_processing/frontal_face_detector.h> // 添加缺失的头文件
#include <vector>
#include <string>
#include <map>

// 网络结构定义
template <template <int,template<typename>class,int,typename> class block, int N, template<typename>class BN, typename SUBNET>
using residual = dlib::add_prev1<block<N,BN,1,dlib::tag1<SUBNET>>>;

template <template <int,template<typename>class,int,typename> class block, int N, template<typename>class BN, typename SUBNET>
using residual_down = dlib::add_prev2<dlib::avg_pool<2,2,2,2,dlib::skip1<dlib::tag2<block<N,BN,2,dlib::tag1<SUBNET>>>>>>;

template <int N, template <typename> class BN, int stride, typename SUBNET> 
using block  = BN<dlib::con<N,3,3,1,1,dlib::relu<BN<dlib::con<N,3,3,stride,stride,SUBNET>>>>>;

template <int N, typename SUBNET> using ares      = dlib::relu<residual<block,N,dlib::affine,SUBNET>>;
template <int N, typename SUBNET> using ares_down = dlib::relu<residual_down<block,N,dlib::affine,SUBNET>>;

template <typename SUBNET> using alevel0 = ares_down<256,SUBNET>;
template <typename SUBNET> using alevel1 = ares<256,ares<256,ares_down<256,SUBNET>>>;
template <typename SUBNET> using alevel2 = ares<128,ares<128,ares_down<128,SUBNET>>>;
template <typename SUBNET> using alevel3 = ares<64,ares<64,ares<64,ares_down<64,SUBNET>>>>;
template <typename SUBNET> using alevel4 = ares<32,ares<32,ares<32,SUBNET>>>;

using anet_type = dlib::loss_metric<dlib::fc_no_bias<128,dlib::avg_pool_everything<
    alevel0<
    alevel1<
    alevel2<
    alevel3<
    alevel4<
    dlib::max_pool<3,3,2,2,dlib::relu<dlib::affine<dlib::con<32,7,7,2,2,
    dlib::input_rgb_image_sized<150>
    >>>>>>>>>>>>;

class FaceRecognizer {
public:
    struct FaceRecord {
        std::string name;
        dlib::matrix<float,0,1> descriptor;
    };

    FaceRecognizer(const std::map<std::string, std::string>& config);
    ~FaceRecognizer();
    
    bool initialize();
    void printFaceLibInfo() const;
    std::vector<dlib::rectangle> detectFaces(const dlib::matrix<dlib::rgb_pixel>& img);
    dlib::matrix<dlib::rgb_pixel> alignFace(const dlib::matrix<dlib::rgb_pixel>& img, const dlib::rectangle& face);
    std::string recognizeFace(const dlib::matrix<dlib::rgb_pixel>& face_chip);
    double compareFaces(const dlib::matrix<dlib::rgb_pixel>& face1, const dlib::matrix<dlib::rgb_pixel>& face2);
    void saveFaceDescriptor(const std::string& name, const dlib::matrix<dlib::rgb_pixel>& face_chip);
    void updateFaceLib();

private:
    std::vector<FaceRecord> loadFaceLib();
    double calculateDistance(const dlib::matrix<float,0,1>& descriptor1, 
                            const dlib::matrix<float,0,1>& descriptor2) const;
    double calculateCosineSimilarity(const dlib::matrix<float,0,1>& descriptor1, 
                                    const dlib::matrix<float,0,1>& descriptor2) const;
    
    std::map<std::string, std::string> config_;
    dlib::frontal_face_detector detector_;
    dlib::shape_predictor sp_;
    anet_type net_;
    std::vector<FaceRecord> face_lib_;
    double threshold_;
    bool initialized_;
};

#endif // FACERECOGNIZER_H